#include "localization.h"
#include "cpp_log.h"
#include "least_sqr_line_fit.h"

static bool first_rand = true;

static double grand(double x, double sigema) {
    if (first_rand) {
        srand(clock());
        first_rand = false;
    }
    double u = rand() % 100000 / 100000.0;
    double v = rand() % 100000 / 100000.0;
    if (u == 0) u += 0.000001;
    return x + sqrt(-2 * log(u)) * cos(2 * M_PI * v) * sigema;
}

static double grand(double sigema) {
    return grand(0, sigema);
}

void localization::set_k(size_t i, double value) {
    if (i < k.size() && value >= -0.000001) {
        if(value < 0.00000001)
            value = 0;
        k[i] = value;
        LOGD("set_k,optimize_k:%d,  optimize_value:%lf", i, value);
    }
}

localization::localization(vector<Vector3d> uwbs, double high) {
    assert (uwbs.size() > 2);
    assert (high > 0.0);
    assert (high < 3.0);
    for (size_t i = 0; i < uwbs.size(); i++) {
        assert (uwbs[i].z() > -1.0);
        assert (uwbs[i].z() < 10.0);
        this->uwbs.push_back(uwbs[i]);
    }
    this->high = high;
    first_locate = true;
    no_init_pose = true;
    no_sig_threshold = 50.0;

    init_center = Vector3d(0, 0, 0);
//    init_tag = Vector3d(-0.1352, 0.0703, high);
    init_tag = Vector3d(-0.155, -0.024, 1.146);
    init_direction = Vector3d(1, 0, 0);
    matrix.setIdentity();
    matrix_odom.setIdentity();
    tread = 0.3836;

    r = 0.07;
    R = 0.31889;

    vector<double>().swap(k);
    k.push_back(1);
    k.push_back(1);
    k.push_back(20);
    k.push_back(20);
    k.push_back(20);
    k.push_back(0.01);
}

void localization::set_no_sig_threshold(double threshold) {
    assert (threshold > 1.0);
    no_sig_threshold = threshold;
}

Vector3d localization::locate_s (vector<double> distance, Vector3d reference, vector<double> weight)
{
    reference = locate (distance,reference,weight);
    double err = error(distance,reference);
    LOGD ("locate_s error: %lf\n", err);
    if (err>0.05) {
        int count =0;
        for (size_t i=0; i<distance.size(); i++)
            if (distance[i]<no_sig_threshold)
                count ++;
        if (count<=3)
            return reference;
        else {
            double min_error = 99999.0;
            for (size_t i=0; i<distance.size(); i++)
                if (distance[i]<no_sig_threshold) {
                    vector<double> dis = distance;
                    dis[i] = 100;
                    Vector3d ref = locate (dis,reference,weight);
                    if (error(dis,ref) < min_error) {
                        min_error = error(dis,ref);
                        reference = ref;
                    }
                }
            LOGD ("locate_s min_error: %lf\n", min_error);
            return reference;
        }
    }
    else
        return reference;
}

double signsqrt(double x) {
    if (x > 0)
        return sqrt(x);
    else
        return -sqrt(-x);
}

Vector3d localization::locate(vector<double> distance, Vector3d reference, vector<double> weight) {
    for (int i = 0; i < 100; i++) {
        Vector3d spring(0, 0, 0);
        size_t sig_count = 0;
        for (size_t i = 0; i < uwbs.size(); i++) {
            if (distance[i] > no_sig_threshold) continue;
            sig_count++;
            if ((reference - uwbs[i]).norm() < 0.0001) continue;
            Vector3d dir = (reference - uwbs[i]).normalized();
            dir *= signsqrt(distance[i] - (reference - uwbs[i]).norm()) * 0.1 * weight[i];
            spring += dir;
        }
        spring /= sig_count;
        reference += spring;
        reference.z() = high;
    }

    double residual = error (distance,reference);
    LOGD ("residual:%lf\n", residual);
    LOGD ("reference,x: %lf  y:%lf  z:%lf\n", reference.x(), reference.y(), reference.z());
    return reference;
}

double sqr(double x) {
    return x * x;
}

double localization::error(vector<double> distance, Vector3d robot) {
    //TODO: no_sig_threshold
    assert (uwbs.size() == distance.size());
    double error = 0.0;
    for (size_t i = 0; i < uwbs.size(); i++) {
        if (distance[i] > no_sig_threshold) continue;
        error += sqr((uwbs[i] - robot).norm() - distance[i]);
    }
    return error;
}

vector<double> localization::calc_std (vector<double> distance, int queue_len)
{
    dis_history.push_back(distance);
    if (dis_history.size()==1) {
        vector<double> v (distance.size(),1);
        return v;
    }
    while (dis_history.size()>queue_len)
        dis_history.erase(dis_history.begin());

    vector<least_sqr_line_fit> fit(distance.size());
    for (size_t i=0; i<distance.size(); i++) {
        vector<double> t;
        for (size_t j=0; j<dis_history.size(); j++)
            t.push_back(dis_history[j][i]);
        fit[i].fit(t);
    }

    vector<double> standard (distance.size(), 0);
    for (size_t i=0; i<distance.size(); i++) {
        for (size_t j=0; j<dis_history.size(); j++) {
            standard[i] += (fit[i].guess(j)-dis_history[j][i]) * (fit[i].guess(j)-dis_history[j][i]);
        }
        standard[i] = sqrt (standard[i]/distance.size());
    }
    double max = 0.0;
    for (size_t i=0; i<standard.size(); i++) {
        if (standard[i]<0.001) standard[i] = 0.001;
        standard[i] = 1/standard[i];
        if (standard[i]>max)
            max = standard[i];
    }
    for (size_t i=0; i<standard.size(); i++)
        standard[i] /= max;
    return standard;
}

Vector3d localization::locate_tag(vector<double> distance) {
    assert (uwbs.size() == distance.size());

    vector<double> weight = calc_std(distance, 10);
    if (first_locate) {
        first_locate = false;
        // looking for the nearest one
        int min_index = -1;
        double min = no_sig_threshold;
        for (size_t i = 0; i < distance.size(); i++)
            if (distance[i] < min) {
                min = distance[i];
                min_index = i;
            }
        assert (min < no_sig_threshold);
        // try 20 times init position randomly around nearest uwb, return the best one
        double min_error = 99999.0;
        Vector3d loc(0, 0, 0);
        for (int i = 0; i < 20; i++) {
            Vector3d reference = uwbs[min_index] + Vector3d(grand(3.0), grand(3.0), 0);
            reference.z() = high;
            Vector3d l = locate_s (distance, reference, weight);
            double e = error(distance, l);
            if (e < min_error) {
                min_error = e;
                loc = l;
            }
        }
        return current_tag = locate_s(distance, loc, weight);
    }
    else {
        return current_tag = locate_s(distance, current_tag, weight);
    }
}

void localization::init_pose(Vector3d center1, Vector3d tag1) {
    assert (center1.z() == 0);

    matrix.setIdentity();
    matrix.translate(center1);
    int num = 1000;
    double min_distance = 99999.0;
    double angle = 0.0;
    for (int i = 0; i < num; i++) {
        Transform<double, 3, Affine> r;
        r.setIdentity();
        r.rotate(AngleAxisd(2 * M_PI * i / num, Vector3d::UnitZ()));
        Vector3d tag_t = matrix * r * init_tag;
        if ((tag_t - tag1).norm() < min_distance) {
            min_distance = (tag_t - tag1).norm();
            angle = 2 * M_PI * i / num;
        }
    }
    matrix.rotate(AngleAxisd(angle, Vector3d::UnitZ()));
    no_init_pose = false;
}

Vector3d localization::get_dir() {
    assert (no_init_pose == false);
    return matrix * matrix_odom * init_direction - matrix * matrix_odom * init_center;
}

Vector3d localization::get_tag() {
    assert (no_init_pose == false);
    return matrix * matrix_odom * init_tag;
}

Vector3d localization::get_center() {
    assert (no_init_pose == false);
    return matrix * matrix_odom * init_center;
}

Transform<double, 3, Affine> localization::two_wheel_robot_move(double lengthL, double lengthR) {
    double theta = (lengthR - lengthL) / tread;
    if (fabs(theta) < 0.0001) {
        Transform<double, 3, Affine> t;
        t.setIdentity();
        t.translate(Vector3d(lengthL, 0, 0));
        return t;
    }
    else {
        double arc = (lengthR + lengthL) / 2;
        double radius = arc / theta;
        Vector3d center = Vector3d(0, radius, 0);

        Transform<double, 3, Affine> t;
        t.setIdentity();
        t.translate(center);
        t.rotate(AngleAxisd(theta, Vector3d::UnitZ()));
        t.translate(-center);
        return t;
    }
}

double localization::norm_2(vector<double> v, vector<double> k) {
    assert (v.size() == k.size());
    double norm = 0.0;
    for (size_t i = 0; i < v.size(); i++)
        norm += v[i] * v[i] * k[i];
    return norm;
}

double localization::d_model(double x, double y, double theta, double dl, double dr, Vector3d tag) {
    Transform<double, 3, Affine> m;
    m.setIdentity();
    m.translate(Vector3d(x, y, 0));
    m.rotate(AngleAxisd(theta, Vector3d::UnitZ()));
    for (size_t i = 0; i < left.size(); i++)
        m = m * two_wheel_robot_move(left[i] * (1 + dl), right[i] * (1 + dr));
    return (matrix * m * init_tag - tag).norm();
}

void localization::locate(vector<double> distance) {
    int count = 0;
    for (size_t i=0; i<distance.size(); i++)
        if (distance[i]<no_sig_threshold)
            count ++;
    if (count<3)
        return;
    locate_tag(distance);
    vector<double> v(6, 0);
    vector<double> gradiant(6);
    v[5] = d_model(0, 0, 0, 0, 0, current_tag);

    for (int i = 0; i < 50; i++) {
        double init_value = norm_2(v, k);
        for (size_t i = 0; i < v.size() - 1; i++) {
            vector<double> v2(v);
            v2[i] += 0.001;
            v2[5] = d_model(v2[0], v2[1], v2[2], v2[3], v2[4], current_tag);
            gradiant[i] = (norm_2(v2, k) - init_value) / 0.001;
            v2[i] -= 0.001;
        }
        for (size_t i = 0; i < v.size() - 1; i++) {
            v[i] -= gradiant[i] / 100.0;
        }
        v[5] = d_model(v[0], v[1], v[2], v[3], v[4], current_tag);
    }

    matrix.translate(Vector3d(v[0], v[1], 0));
    matrix.rotate(AngleAxisd(v[2], Vector3d::UnitZ()));
    size_t queue_size = 50;
    for (size_t i = 0; i < left.size(); i++) {
        left[i] *= 1 + v[3];
        right[i] *= 1 + v[4];
    }
    if (left.size() > queue_size)
        for (size_t i = 0; i < left.size() - queue_size; i++) {
            matrix = matrix * two_wheel_robot_move(left[i], right[i]);
        }
    while (left.size() > queue_size) {
        left.erase(left.begin());
        right.erase(right.begin());
    }
    matrix_odom.setIdentity();
    for (size_t i = 0; i < left.size(); i++) {
        matrix_odom = matrix_odom * two_wheel_robot_move(left[i], right[i]);
    }
}

void localization::set_speed(double left_speed, double right_speed) {
    left.push_back(left_speed);
    right.push_back(right_speed);
    matrix_odom = matrix_odom * two_wheel_robot_move(left_speed, right_speed);
}

void localization::set_speed_angle (double left_speed, double right_speed, double angle, double ratio)
{
    double speed = (left_speed + right_speed)/2.0;
    double theta = (right_speed - left_speed)/tread;
    theta = theta * (1-ratio) + angle * ratio;
    double delta = theta * tread;
    left_speed = speed - delta/2.0;
    right_speed = speed + delta/2.0;

    left.push_back (left_speed);
    right.push_back (right_speed);
    matrix_odom = matrix_odom * two_wheel_robot_move (left_speed,right_speed);
}

void localization::draw_circle(int pice, double r) {
#ifdef QT
    Vector2d p (0,r);
    Rotation2D<double> rot(2*M_PI/pice);
    glBegin (GL_LINE_STRIP);
    for (int i=0; i<pice+1; i++) {
        glVertex3d (p.x(),p.y(),0);
        p = rot * p;
    }
    glEnd ();
#endif
}

void localization::draw(double scale) {
#ifdef QT
    glPushMatrix ();
    glTranslated (current_tag.x(),current_tag.y(),0);
    draw_circle (20,0.02);
    glPopMatrix ();
    if (no_init_pose) return ;

    glPushMatrix ();
    Transform<double,3,Affine> m = matrix * matrix_odom;
    Matrix4d m2 (m.data ());
    Matrix4d m3 = m2.transpose ();
    glMultTransposeMatrixd(m3.data());
    glColor3d (1,0.6,0.6);

    glLineWidth ((int)(sqrt(scale)/5.0));
    draw_circle (50,R);
    draw_circle (20,0.02);
    glBegin (GL_LINES);
        glVertex3d (init_center.x(),init_center.y(),init_center.z());
        Vector3d aim = init_center + init_direction;
        glVertex3d (aim.x()*R,aim.y()*R,aim.z()*R);

        glVertex3d ( r, tread/2,0);
        glVertex3d (-r, tread/2,0);
        glVertex3d ( r,-tread/2,0);
        glVertex3d (-r,-tread/2,0);

        glVertex3d (0, tread/2,0);
        glVertex3d (0,-tread/2,0);
    glEnd ();
    glPopMatrix ();
#endif
}
